#ifndef _TOOLFUNCTION_HPP
#define _TOOLFUNCTION_HPP

#include <vector>
#include <iostream>
#include <fstream>
#include <memory>
#include <cstring>
#include <numeric>
#include <chrono>
#include <algorithm>

#include "ToolFunction.hpp"

bool cmp(Detection &a, Detection &b)
{
    return a.conf > b.conf;
}

float iou(Detection a, Detection b)
{
    float *lbox = a.bbox;
    float *rbox = b.bbox;

    for (int i = 0; i < 4; i++)
    {
        lbox[i] = lbox[i] + a.class_id * 7680;
        rbox[i] = rbox[i] + b.class_id * 7680;
    }

    float interBox[] = {
        std::max(lbox[0], rbox[0]), // left
        std::min(lbox[2], rbox[2]), // right
        std::max(lbox[1], rbox[1]), // top
        std::min(lbox[3], rbox[3]), // bottom
    };

    if (interBox[2] > interBox[3] || interBox[0] > interBox[1])
        return 0.0f;

    float interBoxS = (interBox[1] - interBox[0]) * (interBox[3] - interBox[2]);
    return interBoxS / ((lbox[2] - lbox[0]) * (lbox[3] - lbox[1]) + (rbox[2] - rbox[0]) * (rbox[3] - rbox[1]) - interBoxS + 0.000001f);
}

std::vector<float> letterbox(cv::Mat &src, int h, int w, std::vector<float> &pad)
{
    int INPUT_C = 3;
    int channelLength = h * w;
    int in_w = src.cols;
    int in_h = src.rows;
    int tar_w = w;
    int tar_h = h;
    float r = std::min(float(tar_h) / in_h, float(tar_w) / in_w);
    int inside_w = round(in_w * r);
    int inside_h = round(in_h * r);
    float pad_w = tar_w - inside_w;
    float pad_h = tar_h - inside_h;

    cv::Mat resize_img;

    cv::resize(src, resize_img, cv::Size(inside_w, inside_h));

    float pad_w_1 = pad_w / 2.0;
    float pad_h_1 = pad_h / 2.0;

    pad.push_back(pad_w_1);
    pad.push_back(pad_h_1);
    pad.push_back(r);

    int top = int(round(pad_h_1 - 0.1));
    int bottom = int(round(pad_h_1 + 0.1));
    int left = int(round(pad_w_1 - 0.1));
    int right = int(round(pad_w_1 + 0.1));
    cv::copyMakeBorder(resize_img, resize_img, top, bottom, left, right, 0, cv::Scalar(114, 114, 114));
    cv::cvtColor(resize_img, resize_img, cv::COLOR_BGR2RGB);
    resize_img.convertTo(resize_img, CV_32FC3, 1. / 255.);
    std::vector<cv::Mat> input_channels(INPUT_C);
    cv::split(resize_img, input_channels);
    std::vector<float> result(w * h * INPUT_C);

    auto data = result.data();
    for (int i = 0; i < INPUT_C; ++i)
    {
        memcpy(data, input_channels[i].data, channelLength * sizeof(float));
        data += channelLength;
    }

    return result;
}

std::vector<Detection> nms(float *out_pred, std::vector<float> pad, float nms_thresh, float score_thresh)
{

    std::vector<Detection> dets;
    for (int i = 0; i < 25200; i++)
    {
        float *yolov5_out = out_pred + i * 6;
        if (yolov5_out[4] > score_thresh)
        {
            Detection det;
            memcpy(&det, yolov5_out, sizeof(Detection));
            dets.push_back(det);
        }
    }

    std::sort(dets.begin(), dets.end(), cmp);
    if (dets.size() > 5000)
        dets.erase(dets.begin() + 5000, dets.end());
    for (size_t m = 0; m < dets.size(); ++m)
    {
        auto &item = dets[m];
        for (size_t n = m + 1; n < dets.size(); ++n)
        {
            if (iou(item, dets[n]) > nms_thresh)
            {

                dets.erase(dets.begin() + n);
                --n;
            }
        }
    }
    float ratio_image = pad[5];
    std::vector<Detection> out_dets;
    for (int i = 0; i < dets.size(); i++)
    {
        Detection det;
        float x1 = dets[i].bbox[0];
        float y1 = dets[i].bbox[1];
        float x2 = dets[i].bbox[2];
        float y2 = dets[i].bbox[3];
        int class_id = dets[i].class_id;
        det.bbox[0] = (x1 - pad[3]) / ratio_image;
        det.bbox[1] = (y1 - pad[4]) / ratio_image;
        det.bbox[2] = (x2 - pad[3]) / ratio_image;
        det.bbox[3] = (y2 - pad[4]) / ratio_image;
        det.conf = dets[i].conf;
        det.class_id = dets[i].class_id;
        out_dets.push_back(det);
    }
    return out_dets;
}

#endif
